import os
import time
import pickle
import joblib
import traceback
import pandas as pd
import xgboost as xgb
import lightgbm as lgb
import torch
import torch.nn as nn


# ===================== LSTM临时验证结构(不变)=====================
class TempLSTM(nn.Module):
    def __init__(self, input_size: int = 77, hidden_size: int = 128, 
                 num_layers: int = 2, dropout: float = 0.3):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0.0
        )
        self.layer_norm = nn.LayerNorm(hidden_size)
        self.fc1 = nn.Linear(hidden_size, hidden_size // 2)
        self.fc2 = nn.Linear(hidden_size // 2, 1)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        lstm_out, _ = self.lstm(x)
        last_step_out = lstm_out[:, -1, :]
        norm_out = self.layer_norm(last_step_out)
        fc1_out = self.relu(self.fc1(norm_out))
        drop_out = self.dropout(fc1_out)
        logits = self.fc2(drop_out)
        return self.sigmoid(logits)


def check_pkl_models(model_dir: str) -> pd.DataFrame:
    """
    修复LGB/XGB检测问题: 兼容后缀错误的.pkl(实际是文本格式)和正确后缀.xgb/.lgb
    """
    results = []
    target_suffixes = [".pkl", ".xgb", ".pth"]

    for filename in os.listdir(model_dir):
        file_suffix = os.path.splitext(filename)[1].lower()
        if file_suffix not in target_suffixes:
            continue
        
        model_path = os.path.join(model_dir, filename)
        result = {
            "文件名": filename,
            "路径": model_path,
            "文件大小(KB)": round(os.path.getsize(model_path) / 1024, 2),
            "修改时间": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(os.path.getmtime(model_path))),
            "文件类型": file_suffix[1:].upper(),
            "协议版本/格式": None,
            "可用状态": "未知",
            "加载错误信息": "",
            "模型类型推测": "",
            "备注": ""
        }
        
        try:
            # 1. 小文件过滤(<1KB视为损坏)
            if os.path.getsize(model_path) < 1024:
                result["可用状态"] = "不可用"
                result["加载错误信息"] = "文件过小(<1KB), 疑似损坏"
                results.append(result)
                continue

            # 2. 关键修复: 先判断文件首字符, 区分是否为XGB/LGB文本模型(无论后缀)
            with open(model_path, "rb") as f:
                first_byte = f.read(1)
                # XGB/LGB文本模型首字符是'{'(十六进制7b)
                is_xgb_lgb_text = (first_byte == b'{')

            # 3. 按文件类型+首字符组合处理
            if file_suffix == ".xgb" or (file_suffix == ".pkl" and is_xgb_lgb_text):
                # 处理: 正确后缀.xgb 或 后缀.pkl但实际是XGB文本模型
                result["协议版本/格式"] = "XGB文本协议"
                result["模型类型推测"] = "XGBoost模型(文本格式)"
                try:
                    model = xgb.Booster(model_file=model_path)
                    if model.num_features() > 0:
                        result["可用状态"] = "可用"
                        result["备注"] = f"特征数: {model.num_features()}"
                        # 提示后缀错误(如果是.pkl后缀)
                        if file_suffix == ".pkl":
                            result["备注"] += " | 注意: 文件后缀应为.xgb(当前是.pkl)"
                    else:
                        result["可用状态"] = "不可用"
                        result["加载错误信息"] = "XGB模型无有效特征"
                except Exception as e:
                    result["可用状态"] = "不可用(XGB加载失败)"
                    result["加载错误信息"] = f"XGB文本模型加载失败: {str(e)}"

            elif file_suffix == ".lgb" or (file_suffix == ".pkl" and is_xgb_lgb_text):
                # 处理: 正确后缀.lgb 或 后缀.pkl但实际是LGB文本模型
                result["协议版本/格式"] = "LGB文本协议"
                result["模型类型推测"] = "LightGBM模型(文本格式)"
                try:
                    model = lgb.Booster(model_file=model_path)
                    if model.num_feature() > 0:
                        result["可用状态"] = "可用"
                        result["备注"] = f"特征数: {model.num_feature()}"
                        # 提示后缀错误(如果是.pkl后缀)
                        if file_suffix == ".pkl":
                            result["备注"] += " | 注意: 文件后缀应为.lgb(当前是.pkl)"
                    else:
                        result["可用状态"] = "不可用"
                        result["加载错误信息"] = "LGB模型无有效特征"
                except Exception as e:
                    result["可用状态"] = "不可用(LGB加载失败)"
                    result["加载错误信息"] = f"LGB文本模型加载失败: {str(e)}"

            elif file_suffix == ".pth":
                # LSTM模型检测(不变)
                result["协议版本/格式"] = "PyTorch state_dict(LSTM专用)"
                result["模型类型推测"] = "PyTorch/LSTM模型"
                state_dict = None
                try:
                    state_dict = torch.load(
                        model_path,
                        map_location="cpu",
                        encoding="utf-8",
                        weights_only=False
                    )
                except Exception as e:
                    result["可用状态"] = "不可用(文件异常)"
                    result["加载错误信息"] = f"文件加载失败: {type(e).__name__} | {str(e)}"
                    results.append(result)
                    continue

                if not isinstance(state_dict, dict) or len(state_dict) == 0:
                    result["可用状态"] = "不可用(格式错误)"
                    result["加载错误信息"] = "state_dict为空或非字典格式"
                    results.append(result)
                    continue

                try:
                    temp_model = TempLSTM(
                        input_size=77, hidden_size=128, num_layers=2, dropout=0.3
                    )
                    temp_model.load_state_dict(state_dict)
                    result["可用状态"] = "可用(结构匹配)"
                    result["备注"] = f"权重项数: {len(state_dict)}, 与LSTM结构完全匹配"
                except KeyError as e:
                    result["可用状态"] = "结构不匹配(层名称错误)"
                    result["加载错误信息"] = f"缺失层名称: {str(e)}"
                except RuntimeError as e:
                    result["可用状态"] = "结构不匹配(参数错误)"
                    result["加载错误信息"] = f"维度/参数不匹配: {str(e).splitlines()[0]}"

            elif file_suffix == ".pkl" and not is_xgb_lgb_text:
                # 真正的pickle文件(sklearn/LR/RF)
                result["协议版本/格式"] = None
                result["模型类型推测"] = "sklearn兼容模型"
                with open(model_path, "rb") as f:
                    first_byte = f.read(1)
                    if first_byte == b'(':
                        result["协议版本/格式"] = "0-2 (ASCII)"
                    else:
                        proto_version = ord(first_byte) - 0x80
                        result["协议版本/格式"] = proto_version if 3<=proto_version<=5 else f"未知(首字节: {first_byte.hex()})"
                
                try:
                    model = joblib.load(model_path)
                    result["可用状态"] = "可用"
                    model_class = type(model).__name__
                    if "LogisticRegression" in model_class:
                        result["模型类型推测"] = "逻辑回归(LR)"
                    elif "RandomForest" in model_class:
                        result["模型类型推测"] = "随机森林(RF)"
                    else:
                        result["模型类型推测"] = f"sklearn兼容({model_class})"
                except Exception as e:
                    try:
                        model = pickle.load(open(model_path, "rb"), encoding="utf-8")
                        result["可用状态"] = "可用"
                        result["模型类型推测"] = f"pickle通用({type(model).__name__})"
                    except Exception as e2:
                        result["可用状态"] = "不可用(pickle加载失败)"
                        result["加载错误信息"] = f"pickle文件加载失败: {str(e2)}"

        except Exception as e:
            result["可用状态"] = "不可用(未预料错误)"
            result["加载错误信息"] = f"全局错误: {type(e).__name__} | {str(e)}"
        
        results.append(result)
    
    # 排序整理
    df = pd.DataFrame(results)
    if not df.empty:
        status_order = {
            "可用": 0,
            "可用(结构匹配)": 1,
            "结构不匹配(层名称错误)": 2,
            "结构不匹配(参数错误)": 3,
            "不可用(格式错误)": 4,
            "不可用(文件异常)": 5,
            "不可用(XGB加载失败)": 6,
            "不可用(LGB加载失败)": 7,
            "不可用(pickle加载失败)": 8,
            "不可用(未预料错误)": 9,
            "未知": 10
        }
        df["排序键"] = df["可用状态"].map(status_order).fillna(10)
        df = df.sort_values(by=["排序键", "修改时间"], ascending=[True, False])
        df = df.drop(columns=["排序键"])
    
    return df


# ------------------------------
# 使用示例
# ------------------------------
if __name__ == "__main__":
    MODEL_DIR = "/opt/coding/500w/03.ssq_fractal/models"
    
    print(f"开始检测目录: {MODEL_DIR} 下的模型...")
    report_df = check_pkl_models(MODEL_DIR)
    
    print("\n" + "="*180)
    print(f"模型状态检测报告(共{len(report_df)}个文件)")
    print("="*180)
    pd.set_option('display.max_colwidth', None)
    print(report_df[["文件名", "文件大小(KB)", "可用状态", "加载错误信息", "备注"]].to_string(index=False))
